import mindspore as ms
import pandas as pd
import dgl
from tqdm import tqdm

from ..settings import table

class PerSourceUniform(dgl.dataloading.negative_sampler.PerSourceUniform):
    def _generate(self, g, eids, canonical_etype):
        _, _, vtype = canonical_etype
        shape = eids.shape
        dtype = eids.dtype
        shape = (shape[0] * self.k,)
        src, _ = g.find_edges(eids, etype=canonical_etype)
        src = ms.numpy.repeat(src, self.k, 0)
        dst = ms.numpy.randint(0, g.num_nodes(vtype), shape, dtype=dtype)
        return src, dst


class UniformUnseen(dgl.dataloading.negative_sampler.PerSourceUniform):
    def __init__(self, k, negs, neg_lens):
        super().__init__(k)
        self.negs, self.neg_lens = negs, neg_lens
        self.offset = {ntype: ms.ops.concat([ms.Tensor([0], dtype=ms.int32), neg_len.cumsum(0)[:-1]])
                       for ntype, neg_len in negs.items()}
        # self.offsets = torch.arange(k, device=negs.device).expand_dims(0)

    def _generate(self, g, eids, canonical_etype):
        _, etype, vtype = canonical_etype
        src, _ = g.find_edges(eids, etype=canonical_etype)
        offset = self.offset[vtype][src].expand_dims(-1)
        neg_lens = self.neg_lens[vtype][src]
        dst_idx = ms.numpy.randint(0, int(neg_lens.astype(ms.float16).max()),
                                   (eids.shape[0], self.k))
        dst_idx %= neg_lens.expand_dims(-1)
        dst = self.negs[vtype][dst_idx + offset].view(-1)
        src = ms.numpy.tile(src.expand_dims(-1), (1, self.k)).view(-1)
        # dst = F.randint(shape, dtype, ctx, 0, g.num_nodes(vtype))
        return src, dst


def get_negative_sampler(k, df, num_nodes):
    cols = df.columns
    negs = {}
    neg_lens = {}
    for utype, vtype in [cols, cols[::-1]]:
        all_nodes = set(df[vtype].unique())
        id2pos = {uid: set(group[vtype].unique()) for uid, group in df.groupby(utype)}
        negs[vtype] = []
        neg_lens[vtype] = []
        for uid in tqdm(range(num_nodes[utype])):
            neg = list(all_nodes - id2pos[uid]) if uid in id2pos else list(all_nodes)
            negs[vtype] += neg
            neg_lens[vtype].append(len(neg))
        negs[vtype] = ms.Tensor(negs[vtype], dtype=ms.int32)
        neg_lens[vtype] = ms.Tensor(neg_lens[vtype], dtype=ms.int32)
    return UniformUnseen(k, negs, neg_lens)


def load_interaction_graph(path, data):
    df_train = pd.read_csv(path + 'train.csv', header=None, names=['uid', 'iid', 'timestamp', 'type'])
    df_valid = pd.read_csv(path + 'valid.csv', header=None, names=['uid', 'iid', 'timestamp', 'type'])
    df_test = pd.read_csv(path + 'test.csv', header=None, names=['uid', 'iid', 'timestamp', 'type'])
    edge_dict = {}
    edge_feature_dict = {}
    reverse_etypes = {}
    t_max = df_train['timestamp'].max()
    for k, v in table[data]['edge_name'].items():
        edges = df_train[df_train['type'] == k][['uid', 'iid', 'timestamp']].to_numpy()
        edge_dict[('user', v, 'item')] = (edges[:, 0], edges[:, 1])
        edge_dict[('item', v + '-ed', 'user')] = (edges[:, 1], edges[:, 0])
        edge_feature_dict[('user', v, 'item')] = t_max - ms.Tensor(edges[:, 2])
        edge_feature_dict[('item', v + '-ed', 'user')] = t_max - ms.Tensor(edges[:, 2])
        reverse_etypes[v] = v + '-ed'
        del edges
    return edge_dict, edge_feature_dict


def load_knowledge_graph(path, fname, data, suffix='-ed'):
    srctype, dsttype = fname.split('-')
    if srctype == dsttype:
        suffix = ''
    df = pd.read_csv(f'{path}{srctype}-{dsttype}.csv', header=None, names=['u', 'v', 'type'])
    edge_dict = {}
    for k in df['type'].unique():
        v = table[data]['edge_name'][k] if 'edge_name' in table[data] else k
        edges = df[df['type'] == k].drop(columns=['type']).to_numpy()
        edge_dict[(srctype, v, dsttype)] = (edges[:, 0], edges[:, 1])
        if suffix != '':
            edge_dict[(dsttype, v + suffix, srctype)] = (edges[:, 1], edges[:, 0])
    return srctype, dsttype, df['u'].max(), df['v'].max(), edge_dict


def load_co_relation(df, threshold=0, u='user', v='item', device=None):
    df1 = df[df['type'] > 4][[u, v]]
    df2 = df1.copy()
    df1.columns, df2.columns = ['u1', v], ['u2', v]
    df_merge = df1.merge(df2, on=v).value_counts(['u1', 'u2'])
    u1 = df_merge.index.map(lambda x: x[0])
    u2 = df_merge.index.map(lambda x: x[1])
    df_merge = df_merge[u1 != u2]
    if threshold > 0:
        df_merge = df_merge[df_merge >= threshold]
    edata = {'count': ms.Tensor(df_merge.reset_index(drop=True))}
    edges = (ms.Tensor(df_merge.index.map(lambda x: x[0])),
             ms.Tensor(df_merge.index.map(lambda x: x[1])))
    return edges, edata


def load_graph(df, num_nodes, features,
               global_edge_dict, global_edata_dict,
               t_max=None, suffix='-ed', ntypes=('user', 'item')):
    edge_dict, reverse_etypes = {}, {}
    snapshot_feature_dict = {'timestamp': {}}
    t_max = t_max if t_max else df['timestamp'].max()
    for v, group in df.groupby('type'):
        edges = group.drop(columns=['type']).to_numpy()
        edge_dict[(ntypes[0], v, ntypes[1])] = (edges[:, 0], edges[:, 1])
        snapshot_feature_dict['timestamp'][(ntypes[0], v, ntypes[1])] = t_max - ms.Tensor(edges[:, 2], dtype=ms.int32)
        if ntypes[0] != ntypes[1]:
            edge_dict[(ntypes[1], v + suffix, ntypes[0])] = (edges[:, 1], edges[:, 0])
            snapshot_feature_dict['timestamp'][(ntypes[1], v + suffix, ntypes[0])] = t_max - ms.Tensor(edges[:, 2], dtype=ms.int32)
        # snapshot_feature_dict['timestamp'][('user', v, 'item')] = ms.Tensor(edges[:, 2])
        # snapshot_feature_dict['timestamp'][('item', v, 'user')] = ms.Tensor(edges[:, 2])
        # reverse_etypes[v] = v + suffix
        for etype, edges in global_edge_dict.items():
            edge_dict[etype] = edges
    del edges
    G = dgl.heterograph(edge_dict, num_nodes, idtype=ms.int32, device=ms.context._context())
    for feature_dict in [global_edata_dict, snapshot_feature_dict]:
        for ftype, edge_feature_d in feature_dict.items():
            # for etype, edge_feature in edge_feature_d.items():
            G.edata[ftype] = edge_feature_d
    # G = G.to(device)
    for ntype in G.ntypes:
        G.nodes[ntype].data['features'] = features[ntype]
    print(G)
    return G

def add_edges(old_snapshot, new_snapshot, etypes):
    for etype in etypes:
        edges = new_snapshot.all_edges(form='all', etype=etype)
        old_snapshot.add_edges(edges[0], edges[1], etype=etype,
                                 data=new_snapshot.edges[etype].data if new_snapshot.edges[etype].data else None)
